import os

import requests
from py2neo import Graph

from ai_configs import default_ai_configs

GRAPH_HOST=os.getenv("GRAPH_HOST")
GRAPH_PORT=os.getenv("GRAPH_PORT")
GRAPH_USER=os.getenv("GRAPH_USER")
GRAPH_PASSWORD=os.getenv("GRAPH_PASSWORD")
GRAPH_DATABASE=os.getenv("GRAPH_DATABASE")

graph = Graph(f"bolt://{GRAPH_HOST}:{GRAPH_PORT}", auth=(GRAPH_USER, GRAPH_PASSWORD), name=GRAPH_DATABASE)

_ai_config = default_ai_configs.get("local_embedding")

def get_embedding(text):
    # 设置HTTP请求头，包括内容类型和授权信息
    header = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {_ai_config.get("key")}",
    }
    # 构造请求体，指定模型和输入内容
    payload = {
        "model": _ai_config.get("model"),
        "input": text
    }
    # 发送POST请求嵌入API
    response = requests.post(_ai_config.get("url"), json=payload, headers=header)
    # 如果请求成功
    if response.status_code == 200:
        # 解析返回的JSON数据
        data = response.json()
        # 提取嵌入向量
        embedding = data["data"][0]["embedding"]
        # 返回嵌入向量 2560
        return embedding
    else:
        # 如果请求失败，抛出异常并输出错误信息
        raise Exception(f"Embedding API error: {response.text}")

def query_book_with_embedding(query_embedding, top_k=3):
    query = """
        CALL db.index.vector.queryNodes("book_embeddings", $top_k, $query_embedding)
        YIELD node,score
        MATCH (node:Book)
        OPTIONAL MATCH (node)-[:WRITTEN_BY]->(author:Author)
        OPTIONAL MATCH (node)-[:PUBLISHED_BY]->(publisher:Publisher)
        OPTIONAL MATCH (node)-[:HAS_CATEGORY]->(category:Category)
        OPTIONAL MATCH (node)-[:HAS_KEYWORD]->(keyword:Keyword)
        RETURN node.name as 书名,
        node.publish_year as 出版时间,
        node.summary as 简介,
        score as similarity,
        author.name as 作者,
        publisher.name as 出版社,
        category.name as 类别,
        collect(DISTINCT keyword.name) as 关键字
        ORDER BY score DESC
    """
    results = graph.run(query, query_embedding=query_embedding, top_k=top_k)
    books = []
    for record in results:
        book_info = {
            "name": record["书名"],
            "出版时间": record["出版时间"],
            "简介": record["简介"],
            "similarity": record["similarity"],
            "作者": record["作者"],
            "出版社": record["出版社"],
            "类别": record["类别"],
            "关键字": record["关键字"],
        }
        books.append(book_info)
    return books

def query_author_with_embedding(query_embedding, top_k=3):
    query = """
        CALL db.index.vector.queryNodes("author_embeddings", $top_k, $query_embedding)
        YIELD node,score
        MATCH (node:Author)
        OPTIONAL MATCH (node)<-[:WRITTEN_BY]->(book:Book)
        RETURN node.name as 姓名,
        score as similarity,
        COLLECT(DISTINCT book.name) as 所著图书
        ORDER BY score DESC
    """
    results = graph.run(query, query_embedding=query_embedding, top_k=top_k)
    authors = []
    for record in results:
        author = {
            "name": record["姓名"],
            "similarity": record["similarity"],
            "所著图书": record["所著图书"],
        }
        authors.append(author)
    return authors







